#! /usr/bin/env python
# coding=utf-8

# 将 aa比对 映射到 对应的nt中

import os
import threading
import shutil
import sys
import argparse
from Bio import SeqIO



parser = argparse.ArgumentParser(
    description='''将核酸序列 根据其氨基酸的比对 按 codon进行排序
    用法:
    alignsingle_aa2nt.py -a ali.fa -n new_mRNA.fa -o ali_mRNA.fa
    由大天才于2021年7月20日创建于浙江农业大学''')

parser.add_argument('-a',
                help='必须给定，比对好的氨基酸的文件')


parser.add_argument('-o',
                help='必须给定，输出文件')

parser.add_argument('-n',
                help='必须给定，没比对的核酸的文件')



args = parser.parse_args()

if not args.a or not args.o or not args.n:
    parser.print_help()
    sys.exit()



infile_aa = args.a

outfile_nucl  = args.o

infile_nucl = args.n


outfile = open(outfile_nucl,'w')
nucl_dic = {}
for j in SeqIO.parse(infile_nucl,'fasta'):
	nucl_dic[str(j.name)] = str(j.seq)

# 读取比对序列
for j in SeqIO.parse(infile_aa,'fasta'):
	if str(j.name) in nucl_dic:
		matched_nucl_seq = nucl_dic[str(j.name)]
	else:
		print(j.name)
		continue

	new_seq = ''
	start_point = 0
	for k in j.seq:
		if k =='-':
			new_seq+='---'
		else:
			new_seq += matched_nucl_seq[start_point:start_point+3]
			start_point += 3


	if len(matched_nucl_seq) != len(str(j.seq).replace('-',''))*3+3:
		print(len(matched_nucl_seq), len(str(j.seq).replace('-',''))*3+3)
		print('warning, %s not same length' % j.name)


	outfile.write('>'+j.name+'\n')
	outfile.write(new_seq+'\n')

outfile.close()



